@file:OptIn(ExperimentalSerializationApi::class)

package com.lalameow.researchstation.util.serialization

import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.descriptors.SerialDescriptor
import kotlinx.serialization.descriptors.buildClassSerialDescriptor
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import net.minecraft.block.Block
import net.minecraft.entity.EntityType
import net.minecraft.item.Item
import net.minecraft.stat.Stat
import net.minecraft.stat.StatType
import net.minecraft.util.Identifier
import net.minecraft.util.registry.Registry
import net.silkmc.silk.core.serialization.serializers.ResourceLocationSerializer
import kotlin.reflect.KClass

class StatTypeSerializer<T> : KSerializer<StatType<T>> {
    override val descriptor = SerialDescriptor("StatType<*>", ResourceLocationSerializer.descriptor)
    override fun deserialize(decoder: Decoder) =
        Registry.STAT_TYPE[decoder.decodeSerializableValue(ResourceLocationSerializer)] as StatType<T>

    override fun serialize(encoder: Encoder, value: StatType<T>) =
        encoder.encodeSerializableValue(ResourceLocationSerializer, Registry.STAT_TYPE.getId(value)!!)
}

class IdentifierStatSerializer(private val type: StatType<Identifier>) : KSerializer<Stat<Identifier>> {
    override val descriptor = SerialDescriptor("Stat<Identifier>", ResourceLocationSerializer.descriptor)

    override fun deserialize(decoder: Decoder) =
        type.getOrCreateStat(decoder.decodeSerializableValue(ResourceLocationSerializer))!!

    override fun serialize(encoder: Encoder, value: Stat<Identifier>) {
        encoder.encodeSerializableValue(ResourceLocationSerializer, value.value)
    }
}

class BlockStatSerializer(private val type: StatType<Block>) : KSerializer<Stat<Block>> {
    override val descriptor = SerialDescriptor("Stat<Block>", ResourceLocationSerializer.descriptor)

    override fun deserialize(decoder: Decoder) =
        type.getOrCreateStat(Registry.BLOCK[decoder.decodeSerializableValue(ResourceLocationSerializer)])!!

    override fun serialize(encoder: Encoder, value: Stat<Block>) {
        encoder.encodeSerializableValue(ResourceLocationSerializer, Registry.BLOCK.getId(value.value))
    }
}

class ItemStatSerializer(private val type: StatType<Item>) : KSerializer<Stat<Item>> {
    override val descriptor = SerialDescriptor("Stat<Item>", ResourceLocationSerializer.descriptor)

    override fun deserialize(decoder: Decoder) =
        type.getOrCreateStat(Registry.ITEM[decoder.decodeSerializableValue(ResourceLocationSerializer)])!!

    override fun serialize(encoder: Encoder, value: Stat<Item>) {
        encoder.encodeSerializableValue(ResourceLocationSerializer, Registry.ITEM.getId(value.value))
    }
}

class EntityTypeStatSerializer(private val type: StatType<EntityType<*>>) : KSerializer<Stat<EntityType<*>>> {
    override val descriptor = SerialDescriptor("Stat<EntityType<*>>", ResourceLocationSerializer.descriptor)

    override fun deserialize(decoder: Decoder) =
        type.getOrCreateStat(Registry.ENTITY_TYPE[decoder.decodeSerializableValue(ResourceLocationSerializer)])!!

    override fun serialize(encoder: Encoder, value: Stat<EntityType<*>>) {
        encoder.encodeSerializableValue(ResourceLocationSerializer, Registry.ENTITY_TYPE.getId(value.value))
    }
}

object StatSerializer : KSerializer<Stat<*>> {
    private val identifierDescriptor = ResourceLocationSerializer.descriptor
    private val serializerMap = mapOf<KClass<*>, (type: StatType<*>) -> KSerializer<Stat<*>>>(
        Identifier::class to { IdentifierStatSerializer(it as StatType<Identifier>) as KSerializer<Stat<*>> },
        Block::class to { BlockStatSerializer(it as StatType<Block>) as KSerializer<Stat<*>> },
        Item::class to { ItemStatSerializer(it as StatType<Item>) as KSerializer<Stat<*>> },
        EntityType::class to { EntityTypeStatSerializer(it as StatType<EntityType<*>>) as KSerializer<Stat<*>> }
    )
    override val descriptor = buildClassSerialDescriptor("Stat<*>") {
        element("type", identifierDescriptor)
        element("id", identifierDescriptor)
    }

    override fun deserialize(decoder: Decoder) = decoder.beginStructure(descriptor).let {
        val statType =
            Registry.STAT_TYPE[it.decodeSerializableElement(identifierDescriptor, 0, ResourceLocationSerializer)]!!
        val serializer = serializerMap[statType.generic()]!!(statType)
        it.decodeSerializableElement(serializer.descriptor, 1, serializer)
    }

    override fun serialize(encoder: Encoder, value: Stat<*>) {
        encoder.beginStructure(descriptor).let {
            val statType = value.type
            it.encodeSerializableElement(
                identifierDescriptor,
                0,
                ResourceLocationSerializer,
                Registry.STAT_TYPE.getId(statType)!!
            )
            val serializer = serializerMap[statType.generic()]!!(statType)
            it.encodeSerializableElement(serializer.descriptor, 1, serializer, value)
        }
    }
}

fun <T> StatType<T>.generic() = first().value!!::class
